import torch
import torch.nn.functional as F

from PIL import Image

import os
import json
import shutil
import numpy as np
from matplotlib.colors import LinearSegmentedColormap
import matplotlib.pyplot as plt

import torchvision
from torchvision import models
from torchvision import transforms

from captum.attr import IntegratedGradients, DeepLiftShap, DeepLift
from captum.attr import NoiseTunnel
from captum.attr import visualization as viz

from tqdm import tqdm
from model.SupCon import resnet
import torchvision.transforms as T
import data.transforms as MT
from torch.utils.data import Dataset, DataLoader

output_root = './z_captum_output'
os.makedirs(output_root, exist_ok=True)

os.environ['CUDA_VISIBLE_DEVICES'] = '1'
device = torch.device('cuda:0')

load_model = '.details151/MD/12-13_15:48:33_5最优权重正梯度/Net_best.pth'
net = resnet(34, n_channels=2, n_classes=2)
net.load_state_dict(torch.load(load_model, map_location=device))
net.to(device)
net.eval()

# ig = IntegratedGradients(net)
# nt= NoiseTunnel(ig)
dl = DeepLift(net)
dls = DeepLiftShap(net)
baseline_dist = (torch.randn(10, 2,81,81) * 0.001).to(device)

class PairDataset(Dataset):
    def __init__(self, ct, cta, transform):
        self.transform = transform
        self.datas = []
        categories = ['1']#sorted(os.listdir(ct))
        for label, cate in enumerate(categories):
            ct_path = os.path.join(ct, cate)
            cta_path = os.path.join(cta, cate)
            for imgcta in sorted(os.listdir(cta_path)):
                names = imgcta.split('_')
                if len(names) != 4:
                    continue
                imgct = os.path.join(ct_path, '_'.join(names[:2]+names[3:]))
                if os.path.exists(imgct):
                    imgcta = os.path.join(cta_path, imgcta)
                    self.datas.append( (imgct, imgcta, label) )
    
    def __len__(self):
        return len(self.datas)
    
    def __getitem__(self, index):
        ctimgpath, ctaimgpath, label = self.datas[index]
        ctimg = self.transform(Image.open(ctimgpath))
        ctaimg = self.transform(Image.open(ctaimgpath))
        return ctimg, ctaimg, label, os.path.basename(ctaimgpath)

transform = T.Compose([
    T.Resize(81),
    T.CenterCrop(81),
    T.ToTensor(),
    MT.SobelChannel(3)
])
    
dataset = PairDataset('/nfs3-p1/zsxm/adpaper/4visualize/ct', '/nfs3-p1/zsxm/adpaper/4visualize/cta', transform)

for ctimg, ctaimg, label, file in tqdm(dataset):
    # if label != 1:
    #     continue
    
    img = ctimg.unsqueeze(0).to(device)
    
    # attributions_1 = ig.attribute(img, target=1)
    # attributions_2 = nt.attribute(img, nt_samples=10, nt_type='smoothgrad_sq', target=1)
    attributions_1 = dl.attribute(img, target=1)
    attributions_2 = dls.attribute(img, baselines=baseline_dist, target=1)

    plt_fig = plt.figure(figsize=(12,3), dpi=200)
    plt_axis = plt_fig.subplots(1, 4)
    viz.visualize_image_attr(attributions_1[0,0].cpu().detach().unsqueeze(-1).repeat(1,1,3).numpy(),
                            ctimg[0].unsqueeze(-1).repeat(1,1,3).numpy(),
                            method="original_image",
                            sign="all",
                            plt_fig_axis=[plt_fig, plt_axis[0]],
                            show_colorbar=True,
                            outlier_perc=1,
                            use_pyplot=False)
    viz.visualize_image_attr(attributions_1[0,0].cpu().detach().unsqueeze(-1).repeat(1,1,3).numpy(),
                            ctimg[0].unsqueeze(-1).repeat(1,1,3).numpy(),
                            method="heat_map",
                            sign="positive",
                            plt_fig_axis=[plt_fig, plt_axis[1]],
                            show_colorbar=True,
                            outlier_perc=1,
                            use_pyplot=False)
    viz.visualize_image_attr(attributions_1[0,0].cpu().detach().unsqueeze(-1).repeat(1,1,3).numpy(),
                            ctimg[0].unsqueeze(-1).repeat(1,1,3).numpy(),
                            method='blended_heat_map',
                            sign="positive",
                            plt_fig_axis=[plt_fig, plt_axis[2]],
                            show_colorbar=True,
                            outlier_perc=1,
                            alpha_overlay=0.3,
                            use_pyplot=False)
    viz.visualize_image_attr(attributions_1[0,0].cpu().detach().unsqueeze(-1).repeat(1,1,3).numpy(),
                            ctaimg[0].unsqueeze(-1).repeat(1,1,3).numpy(),
                            method="original_image",
                            sign="all",
                            plt_fig_axis=[plt_fig, plt_axis[3]],
                            show_colorbar=True,
                            outlier_perc=1,
                            use_pyplot=False)
    plt_fig.savefig(os.path.join(output_root, f'dl_{file}'), bbox_inches='tight')
    plt.close(plt_fig)

    plt_fig = plt.figure(figsize=(12,3), dpi=200)
    plt_axis = plt_fig.subplots(1, 4)
    viz.visualize_image_attr(attributions_2[0,0].cpu().detach().unsqueeze(-1).repeat(1,1,3).numpy(),
                            ctimg[0].unsqueeze(-1).repeat(1,1,3).numpy(),
                            method="original_image",
                            sign="all",
                            plt_fig_axis=[plt_fig, plt_axis[0]],
                            show_colorbar=True,
                            outlier_perc=1,
                            use_pyplot=False)
    viz.visualize_image_attr(attributions_2[0,0].cpu().detach().unsqueeze(-1).repeat(1,1,3).numpy(),
                            ctimg[0].unsqueeze(-1).repeat(1,1,3).numpy(),
                            method="heat_map",
                            sign="positive",
                            plt_fig_axis=[plt_fig, plt_axis[1]],
                            show_colorbar=True,
                            outlier_perc=1,
                            use_pyplot=False)
    viz.visualize_image_attr(attributions_2[0,0].cpu().detach().unsqueeze(-1).repeat(1,1,3).numpy(),
                            ctimg[0].unsqueeze(-1).repeat(1,1,3).numpy(),
                            method='blended_heat_map',
                            sign="positive",
                            plt_fig_axis=[plt_fig, plt_axis[2]],
                            show_colorbar=True,
                            outlier_perc=1,
                            alpha_overlay=0.3,
                            use_pyplot=False)
    viz.visualize_image_attr(attributions_2[0,0].cpu().detach().unsqueeze(-1).repeat(1,1,3).numpy(),
                            ctaimg[0].unsqueeze(-1).repeat(1,1,3).numpy(),
                            method="original_image",
                            sign="all",
                            plt_fig_axis=[plt_fig, plt_axis[3]],
                            show_colorbar=True,
                            outlier_perc=1,
                            use_pyplot=False)
    plt_fig.savefig(os.path.join(output_root, f'dls_{file}'), bbox_inches='tight')
    plt.close(plt_fig)
    del img